import json
import logging
import threading
import time
from typing import Dict, Generator, Type, TypeVar, TypedDict
from jinja2 import Environment, FileSystemLoader

from ..configs import TEMPLATE_FOLDER
from .utils import size_bytes_to_human_readable

logger = logging.getLogger(__file__)


class ProgressDetail(TypedDict):
    current: int
    total: int


class DockerPullLog(TypedDict):
    status: str
    progressDetail: ProgressDetail
    progress: str
    id: str


class DockerPullStatus(TypedDict):
    status: str
    total_bytes: int
    current_bytes: int
    progress: str
    id: str


T = TypeVar("T", bound=dict)


class PullingLogManager:
    def __init__(self, stream, blocking_readstream=False):
        self.stream = stream
        self.pulling_target: str = ""
        self.pull_status: Dict[str, DockerPullStatus] = {}
        self.already_exists = False
        if blocking_readstream:
            self.log_reading_thread = threading.Thread(target=self.process_json_stream)
            self.log_reading_thread.daemon = True
            self.log_reading_thread.start()

    def render_pulling_progress_html(self):
        env = Environment(loader=FileSystemLoader(TEMPLATE_FOLDER))
        # 注册过滤器
        env.filters["size_bytes_to_human_readable"] = size_bytes_to_human_readable
        template = env.get_template("components/docker_pull_progress_bar.html")
        return template.render(
            pull_status=self.pull_status,
            pulling_target=self.pulling_target,
            already_exists=self.already_exists,
        )

    def get_log_html(self) -> str:
        """
        Get the formatted html log
        """
        return self.render_pulling_progress_html()

    def process_json_stream(self) -> None:
        chunk: bytes
        for chunk in self.stream:
            line = chunk.decode("utf-8", errors="ignore").strip()
            print(line)
            if line:
                try:
                    log: DockerPullLog = json.loads(line)
                    if log.get("status", "").find("Pulling from") != -1:
                        self.pulling_target = log["status"] + log["id"]
                    else:
                        if "id" in log:
                            # 填充 pull_status 数据
                            log_id = log["id"]
                            progress_detail = log.get("progressDetail", {})
                            current_bytes = progress_detail.get(
                                "current",
                                self.pull_status.get(log_id, {}).get(
                                    "current_bytes", 0
                                ),
                            )
                            total_bytes = progress_detail.get(
                                "total",
                                self.pull_status.get(log_id, {}).get("total_bytes", 0),
                            )
                            progress = log.get("progress", "")
                            status = log.get("status", "")
                            if status == "Download complete":
                                current_bytes = total_bytes

                            self.pull_status[log_id] = {
                                "status": status,
                                "total_bytes": total_bytes,
                                "current_bytes": current_bytes,
                                "progress": progress,
                                "id": log_id,
                            }
                        elif log.get("status", "").find("Image is up to date") != -1:
                            self.already_exists = True
                    # yield log
                except json.JSONDecodeError:
                    print(line)
        # Render HTML after processing all logs
        # html_output = self.render_pulling_progress_html(logs)


if __name__ == "__main__":
    from docker import from_env, DockerClient

    client: DockerClient = from_env()
    stream = client.api.pull(
        # "docker-registry.sky-group.cn/skygroup-library/development-servers:latex-latest-with-cjk",
        "docker-registry.sky-group.cn/skygroup-library/development-servers:ubuntu-22.04",
        stream=True,
    )
    log_mgr = PullingLogManager(stream, blocking_readstream=True)
    time.sleep(1)
    with open("progress.html", "w") as f:
        f.write(log_mgr.get_log_html())
